/*
 * Copyright 2022 The TensorFlow Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *             http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.google.mediapipe.examples.handlandmarker

import android.annotation.SuppressLint
import android.content.Context
import android.content.res.AssetManager
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Paint
import android.util.AttributeSet
import android.util.Log
import android.view.View
import androidx.core.content.ContextCompat
import com.google.mediapipe.tasks.vision.core.RunningMode
import com.google.mediapipe.tasks.vision.handlandmarker.HandLandmarker
import com.google.mediapipe.tasks.vision.handlandmarker.HandLandmarkerResult
import kotlin.math.max
import kotlin.math.min
import kotlin.math.sqrt
import com.google.mediapipe.tasks.components.containers.NormalizedLandmark
import org.tensorflow.lite.Interpreter
import java.io.BufferedReader
import java.io.FileInputStream
import java.io.InputStreamReader
import kotlin.math.pow
import java.nio.channels.FileChannel
import java.nio.ByteBuffer
import kotlin.math.exp

class OverlayView(context: Context?, attrs: AttributeSet?) :
    View(context, attrs) {

    private var results: HandLandmarkerResult? = null
    private var linePaint = Paint()
    private var pointPaint = Paint()
    private val textPaint = Paint()
    private var scaleFactor: Float = 1f
    private var imageWidth: Int = 1
    private var imageHeight: Int = 1
    private val threshold: Float = 0.95f
    private var modelFile= "model.tflite"
    private var trueLabel: String = ""
    private var prob: Float = 0.0f
    private var interpreter: Interpreter? = null
    private var labels: Map<Int, String>? = null
    private var result = Array(1) { FloatArray(2) }
    init {
        initPaints()
        initializeInterpreter()
        loadLabels()
    }

    fun clear() {
        results = null
        linePaint.reset()
        pointPaint.reset()
        textPaint.reset()
        invalidate()
        initPaints()
    }
    private fun loadLabels(){
        val labels = mutableMapOf<Int, String>()

        try {
            context.assets.open("label.txt").use { inputStream ->
                BufferedReader(InputStreamReader(inputStream)).use { reader ->
                    var line: String?
                    while (reader.readLine().also { line = it } != null) {
                        val parts = line?.split(",")
                        if (parts != null && parts.size == 2) {
                            val key = parts[0].trim().toIntOrNull()
                            val value = parts[1].trim()
                            if (key != null && value.isNotEmpty()) {
                                labels[key] = value
                            } else {
                                Log.e("LabelLoader", "无法解析行: $line")
                            }
                        } else {
                            Log.e("LabelLoader", "行格式不正确: $line")
                        }
                    }
                }
            }
        } catch (e: Exception) {
            Log.e("LabelLoader", "读取 label.txt 文件时出错", e)
        }
        Log.d("labels", labels.toString())
        this.labels = labels
    }
    private fun initializeInterpreter() {
        // Load the TF Lite model
        val assetManager = context.assets
        val model = loadModelFile(assetManager)

        // Initialize TF Lite Interpreter with NNAPI enabled
        val options = Interpreter.Options()
        options.setUseNNAPI(true)

        val interpreter = Interpreter(model, options)
        // Read input shape from model file
        val inputShape = interpreter.getInputTensor(0).shape()
        // Read output shape from model file
        val outputShape = interpreter.getOutputTensor(0).shape()
        Log.d("input", "Input Shape: ${inputShape.joinToString(", ")}")
        Log.d("output", "Output Shape: ${outputShape.joinToString(", ")}")
        // Finish interpreter initialization
        this.interpreter = interpreter
        Log.d("tfLite", "init successfully!")
    }
    private fun loadModelFile(assetManager: AssetManager): ByteBuffer {
        val fileDescriptor = assetManager.openFd(modelFile)
        val inputStream = FileInputStream(fileDescriptor.fileDescriptor)
        val fileChannel = inputStream.channel
        val startOffset = fileDescriptor.startOffset
        val declaredLength = fileDescriptor.declaredLength
        return fileChannel.map(FileChannel.MapMode.READ_ONLY, startOffset, declaredLength)
    }

    private fun initPaints() {
        linePaint.color =
            ContextCompat.getColor(context!!, R.color.bottom_sheet_background)
        linePaint.strokeWidth = LINE_STROKE_WIDTH
        linePaint.style = Paint.Style.STROKE
        textPaint.color = Color.RED
        textPaint.textSize = 38f
        textPaint.textAlign = Paint.Align.LEFT
        pointPaint.color = Color.RED
        pointPaint.strokeWidth = POINT_STROKE_WIDTH
        pointPaint.style = Paint.Style.FILL
    }
    private fun calculateDistance(p1: NormalizedLandmark, p2: NormalizedLandmark, scaledX: Float, scaledY: Float): Float {
        // 计算两个点之间的欧几里得距离
        return sqrt(((p1.x() - p2.x()) / scaledX).pow(2) + ((p1.y() - p2.y()) / scaledY).pow(2) + (p1.z() - p2.z()).pow(2))
    }

    private fun calculateRectangle(handLandmarks: MutableList<NormalizedLandmark>): FloatArray {
        val minX = handLandmarks.minOf { it.x() }
        val minY = handLandmarks.minOf { it.y() }
        val maxX = handLandmarks.maxOf { it.x() }
        val maxY = handLandmarks.maxOf { it.y() }
        return floatArrayOf(minX, minY, maxX, maxY)
    }
    private fun getAllDistances(handLandmarks: MutableList<NormalizedLandmark>, minX: Float, minY: Float, maxX: Float, maxY: Float): FloatArray {
        // 计算 handLandmarks 中所有点两两之间的距离

        val scaledX = maxX - minX
        val scaledY = maxY - minY
        val distances = mutableListOf<Float>()
        val numLandmarks = handLandmarks.size
        for (i in 0 until numLandmarks) {
            for (j in i + 1 until numLandmarks) {
                val distance = calculateDistance(handLandmarks[i], handLandmarks[j], scaledX, scaledY)
                distances.add(distance)
            }
        }
        return distances.toFloatArray()
    }
    private fun printFloatArray(floatArray: FloatArray) {
        // 使用 joinToString 将 FloatArray 转换为逗号分隔的字符串
        val floatArrayString = floatArray.joinToString(", ")

        // 使用 Log.d 打印字符串
        Log.d("FloatArray", "FloatArray: $floatArrayString")
    }
    private fun softmax(floatArray: FloatArray): FloatArray {
        // Step 1: Calculate the exponential of each element
        val expValues = FloatArray(floatArray.size)
        var maxVal = Float.NEGATIVE_INFINITY
        for (value in floatArray) {
            if (value > maxVal) {
                maxVal = value
            }
        }

        var sumExp = 0.0
        for (i in floatArray.indices) {
            expValues[i] = exp((floatArray[i] - maxVal).toDouble()).toFloat()
            sumExp += expValues[i]
        }

        // Step 2: Normalize by dividing each exponential value by the sum of all exponentials
        for (i in expValues.indices) {
            expValues[i] /= sumExp.toFloat()
        }

        return expValues
    }
    private fun findMaxValueAndIndex(floatArray: FloatArray): Pair<Float, Int> {
        if (floatArray.isEmpty()) {
            throw IllegalArgumentException("数组不能为空")
        }

        var maxIndex = 0
        var maxValue = floatArray[0]

        for (i in floatArray.indices) {
            if (floatArray[i] > maxValue) {
                maxValue = floatArray[i]
                maxIndex = i
            }
        }

        return Pair(maxValue, maxIndex)
    }
    @SuppressLint("DefaultLocale")
    override fun draw(canvas: Canvas) {
        super.draw(canvas)
        results?.let { handLandmarkResult ->
            for (landmark in handLandmarkResult.landmarks()) {
                val (minX, minY, maxX, maxY) = calculateRectangle(landmark)
                val distances = getAllDistances(landmark, minX, minY, maxX, maxY)
                // 运行推理
                try {
                    interpreter?.run(distances, result)
                } catch (e: Exception) {
                    Log.e("inference", "Error during inference: ${e.message}", e)
                }
                val expValue = softmax(result[0])
                val result = findMaxValueAndIndex(expValue)
//                Log.d("MaxValueFinder", "最大值是 ${result.first}，其下标是 ${result.second}")
                if (result.first >= threshold){
                    trueLabel = labels?.get(result.second).toString()
                    prob = result.first
                }
                else{
                    trueLabel = "none"
                    prob = result.first
                }
                // 计算缩放后的坐标
                val scaledMinX = minX * imageWidth * scaleFactor
                val scaledMinY = minY * imageHeight * scaleFactor
                val scaledMaxX = maxX * imageWidth * scaleFactor
                val scaledMaxY = maxY * imageHeight * scaleFactor

                // 定义矩形的四条边的坐标
                val lines = floatArrayOf(
                    scaledMinX, scaledMinY, scaledMaxX, scaledMinY, // 上边
                    scaledMaxX, scaledMinY, scaledMaxX, scaledMaxY, // 右边
                    scaledMaxX, scaledMaxY, scaledMinX, scaledMaxY, // 下边
                    scaledMinX, scaledMaxY, scaledMinX, scaledMinY  // 左边
                )
                canvas.drawLines(lines, linePaint)
                val formattedProb = String.format("%.2f", prob * 100)
                canvas.drawText("标签${trueLabel}, 概率${formattedProb}%", scaledMinX + 10, scaledMinY - 10, textPaint)
                for (normalizedLandmark in landmark) {
                    canvas.drawPoint(
                        normalizedLandmark.x() * imageWidth * scaleFactor,
                        normalizedLandmark.y() * imageHeight * scaleFactor,
                        pointPaint
                    )
                }

                HandLandmarker.HAND_CONNECTIONS.forEach {
                    canvas.drawLine(
                        landmark[it!!.start()]
                            .x() * imageWidth * scaleFactor,
                        landmark[it.start()]
                            .y() * imageHeight * scaleFactor,
                        landmark[it.end()]
                            .x() * imageWidth * scaleFactor,
                        landmark[it.end()]
                            .y() * imageHeight * scaleFactor,
                        linePaint
                    )
                }
            }
        }
    }

    fun setResults(
        handLandmarkResults: HandLandmarkerResult,
        imageHeight: Int,
        imageWidth: Int,
        runningMode: RunningMode = RunningMode.IMAGE
    ) {
        results = handLandmarkResults

        this.imageHeight = imageHeight
        this.imageWidth = imageWidth

        scaleFactor = when (runningMode) {
            RunningMode.IMAGE,
            RunningMode.VIDEO -> {
                min(width * 1f / imageWidth, height * 1f / imageHeight)
            }
            RunningMode.LIVE_STREAM -> {
                // PreviewView is in FILL_START mode. So we need to scale up the
                // landmarks to match with the size that the captured images will be
                // displayed.
                max(width * 1f / imageWidth, height * 1f / imageHeight)
            }
        }
        invalidate()
    }

    companion object {
        private const val LINE_STROKE_WIDTH = 8F
        private const val POINT_STROKE_WIDTH = 14F
    }
}
